"""统一扩展基类和匹配规则系统

提供统一的扩展处理器基类和灵活的匹配规则系统，支持装饰器处理、内置函数转换等各种扩展功能。
"""

from abc import ABC, abstractmethod
from typing import Optional, Dict, List, Union, Callable, Any

from mcpy.ir.base import IRNode
from mcpy.passes.base import PassStage

from .ir import IRContext
from .ir import IRModule


def condition_to_str(condition: Dict[str, str]) -> str:
    """将条件转换为字符串"""
    return f"{', '.join([f'{k}={v}' for k, v in condition.items()])}"


class MatchRule:
    """匹配规则定义

    支持字典和列表语法：
    - 字典形式：{"decorator": "main", "node_type": "FunctionDef"} (AND关系)
    - 列表形式：[{"decorator": "main"}, {"decorator": "export"}] (OR关系)
    """

    conditions: List[Dict[str, str]]

    def __init__(
        self,
        conditions: Union[Dict[str, str], List[Dict[str, str]]],
        priority: int = 100,
    ):
        if isinstance(conditions, dict):
            self.conditions = [conditions]
        else:
            self.conditions = conditions
        self.priority = priority

    def matches(self, context: IRContext) -> bool:
        """检查目标是否匹配此规则"""
        return any(self._match_condition(cond, context) for cond in self.conditions)

    def _match_condition(
        self,
        condition: Dict[str, str],
        context: IRContext,
    ) -> bool:
        """匹配单个条件"""
        for key, value in condition.items():
            if key == "decorator":
                if not context.get_decorator(value):
                    return False
            elif key == "node_type":
                if context.get_node_type() != value:
                    return False
            elif key == "function_call":
                if context.has_function_call(value):
                    return False
            else:
                if context.get(key) != value:
                    return False
        return True

    def __str__(self) -> str:
        if len(self.conditions) == 0:
            return ""
        elif len(self.conditions) == 1:
            return condition_to_str(self.conditions[0])
        else:
            return (
                f"[{', '.join([condition_to_str(cond) for cond in self.conditions])}]"
            )


class ExtensionProcessor(ABC):
    """统一扩展处理器基类

    所有扩展处理器（装饰器处理器、内置函数处理器等）都应继承此类。
    """

    def __init__(self, name: str):
        self.name = name
        self._match_rules: List[MatchRule] = []
        # 只有在子类没有设置时才使用默认值
        if not hasattr(self, "stage"):
            self.stage: str = "post_analysis"  # 默认阶段
        if not hasattr(self, "timing"):
            self.timing: str = "post"  # 默认时机
        self._setup_rules()
        self._register_convenience_methods()

    @abstractmethod
    def _setup_rules(self) -> None:
        """设置匹配规则

        子类应在此方法中调用 add_rule() 来注册匹配规则。
        """
        pass

    def _register_convenience_methods(self) -> None:
        """注册便利方法

        子类可以重写这个方法来为IRNode注册自己的便利方法。
        """
        pass

    def register_node_method(self, method_name: str, method_func: Callable) -> None:
        """为IRNode注册便利方法

        Args:
            method_name: 方法名
            method_func: 方法函数，第一个参数应该是IRNode实例
        """
        # 动态添加方法到IRNode类
        from .ir import IRNode

        setattr(IRNode, method_name, method_func)

    def add_rule(self, rule: MatchRule) -> None:
        """添加匹配规则"""
        self._match_rules.append(rule)

    def get_rules(self) -> List[MatchRule]:
        """获取所有匹配规则"""
        return self._match_rules.copy()

    def process(self, context: IRContext) -> None:
        """处理IR节点

        Args:
            node: 要处理的IR节点
            context: 处理上下文

        Returns:
            处理状态，决定是否继续处理后续扩展
        """
        pass

    def finally_process(self, context: IRContext) -> None:
        """模块级后处理

        在所有节点处理完成后调用，用于模块级的后处理操作。

        Args:
            ir_module: IR模块
            context: 处理上下文
        """
        pass


class ExtensionRegistry:
    """扩展处理器注册表"""

    def __init__(self):
        self._processors: List[ExtensionProcessor] = []

        # 简化的索引结构：按条件类型快速查找
        self._decorator_index: Dict[str, List[ExtensionProcessor]] = (
            {}
        )  # 装饰器名 -> 处理器列表
        self._node_type_index: Dict[str, List[ExtensionProcessor]] = (
            {}
        )  # 节点类型 -> 处理器列表

        # 阶段化索引：按阶段和时机组织扩展
        self._stage_index: Dict[str, Dict[str, List[ExtensionProcessor]]] = {
            "pre_analysis": {"pre": [], "post": []},
            "type_analysis": {"pre": [], "post": []},
            "post_analysis": {"pre": [], "post": []},
            "optimization": {"pre": [], "post": []},
            "code_gen": {"pre": [], "post": []},
        }

    def register(self, processor: ExtensionProcessor) -> None:
        """注册扩展处理器"""
        self._processors.append(processor)

        # 为新的简化规则构建索引
        for rule in processor.get_rules():
            self._build_index_for_rule(rule, processor)

        # 按优先级排序
        self._sort_processors()

    def _build_index_for_rule(
        self, rule: MatchRule, processor: ExtensionProcessor
    ) -> None:
        """为规则构建索引

        Args:
            rule: 匹配规则
            processor: 扩展处理器
        """
        for condition in rule.conditions:
            if "decorator" in condition:
                decorator_name = condition["decorator"]
                if decorator_name not in self._decorator_index:
                    self._decorator_index[decorator_name] = []
                if processor not in self._decorator_index[decorator_name]:
                    self._decorator_index[decorator_name].append(processor)

            # 为节点类型条件建立索引
            if "node_type" in condition:
                node_type = condition["node_type"]
                if node_type not in self._node_type_index:
                    self._node_type_index[node_type] = []
                if processor not in self._node_type_index[node_type]:
                    self._node_type_index[node_type].append(processor)

    def _sort_processors(self) -> None:
        """按优先级排序处理器"""

        def get_min_priority(processor: ExtensionProcessor) -> int:
            rules = processor.get_rules()
            return min(rule.priority for rule in rules) if rules else 100

        self._processors.sort(key=get_min_priority)

        # 对索引中的处理器也进行排序
        for processors in self._decorator_index.values():
            processors.sort(key=get_min_priority)
        for processors in self._node_type_index.values():
            processors.sort(key=get_min_priority)

    def find_processors(self, context: IRContext) -> List[ExtensionProcessor]:
        """查找可以处理给定IR节点的处理器

        Args:
            node: IR节点
            context: 处理上下文

        Returns:
            可以处理该节点的处理器列表（按优先级排序）
        """
        # 使用索引快速获取候选处理器
        candidate_processors = set()
        node_type = context.get_node_type()
        if node_type and node_type in self._node_type_index:
            # 在索引查找时就过滤 stage 和 timing
            for processor in self._node_type_index[node_type]:
                if (
                    processor.stage == context.stage
                    and processor.timing == context.timing
                ):
                    candidate_processors.add(processor)

        # 根据装饰器获取候选处理器
        for decorator_name in context.get_decorator_names():
            if decorator_name and decorator_name in self._decorator_index:
                # 在索引查找时就过滤 stage 和 timing
                for processor in self._decorator_index[decorator_name]:
                    if (
                        processor.stage == context.stage
                        and processor.timing == context.timing
                    ):
                        candidate_processors.add(processor)

        if len(candidate_processors) == 0:
            return []

        # 对候选处理器进行精确匹配
        matching_processors = []
        for processor in candidate_processors:
            if self._matches_rules(processor, context):
                matching_processors.append(processor)

        # 按优先级排序
        matching_processors.sort(key=self._get_processor_priority)

        return matching_processors

    def _matches_rules(
        self,
        processor: ExtensionProcessor,
        context: IRContext,
    ) -> bool:
        """检查处理器的规则是否匹配

        Args:
            processor: 扩展处理器
            node: IR节点
            context: 处理上下文

        Returns:
            如果规则匹配则返回True，否则返回False
        """
        rules = processor.get_rules()
        if not rules:
            return False

        # 所有规则都必须匹配
        for rule in rules:
            if not rule.matches(context):
                return False

        return True

    def _get_processor_priority(self, processor: ExtensionProcessor) -> int:
        """获取处理器的最小优先级

        Args:
            processor: 扩展处理器

        Returns:
            处理器的最小优先级值
        """
        rules = processor.get_rules()
        return min(rule.priority for rule in rules) if rules else 100

    def get_all_processors(self) -> List[ExtensionProcessor]:
        """获取所有注册的处理器"""
        return self._processors.copy()

    def run_extensions(
        self, stage: PassStage, timing: str, ir_module: IRModule, compiler=None
    ) -> Any:
        """运行 extensions

        Args:
            stage: 执行阶段
            timing: 执行时机
            ir_module: IR 模块
            compiler: 编译器实例（可选）
        """

        context = IRContext(
            module=ir_module,
            stage=stage.value,
            timing=timing,
            compiler=compiler,
        )
        # 只遍历当前模块的直接子节点，不遍历导入模块中的节点
        # 这样可以避免处理导入模块中的函数节点
        for node in ir_module.nodes:
            self._walk_node_recursive(node, context)
        
        # 遍历模块级别的注解（这些注解不在 nodes 中，而是在 annotations 中）
        # 模块级别的注解（如 @mcpy_headers, @mcpy_namespace）需要单独处理
        for annotation in ir_module.annotations:
            self._walk_node_recursive(annotation, context)

        return ir_module

    def _walk_node_recursive(self, node: IRNode, context: IRContext):
        """递归遍历节点及其子节点"""
        context.node = node
        processors = self.find_processors(context)
        if processors:
            for processor in processors:
                processor.process(context)

        # 递归遍历子节点
        for child in node.get_child_nodes():
            self._walk_node_recursive(child, context)


# 全局注册表实例
_global_registry = ExtensionRegistry()


def get_global_extension_registry() -> ExtensionRegistry:
    """获取全局扩展注册表"""
    return _global_registry


def register_extension(processor: ExtensionProcessor) -> None:
    """注册扩展处理器到全局注册表"""
    _global_registry.register(processor)


def extension_processor(
    name: str,
    match_rule: Union[Dict[str, str], List[Dict[str, str]]],
    priority: int = 100,
    stage: str = "post_analysis",
    timing: str = "post",
):
    """扩展处理器装饰器

    自动生成ExtensionProcessor子类，简化扩展的编写。

    Args:
        name: 扩展处理器名称
        match_rule: 匹配规则，支持字典（AND关系）或列表（OR关系）
        priority: 处理优先级，数值越小优先级越高
        stage: 执行阶段 (pre_analysis, type_analysis, post_analysis, optimization, code_gen)
        timing: 执行时机 ("pre" 或 "post")

    Returns:
        装饰后的函数会被包装成ExtensionProcessor子类

    Example:
        # 默认在 post_analysis 阶段的 post 时机执行
        @extension_processor(
            "function_call",
            match_rule=[{"node_type": "call"}],
            priority=1
        )
        def process_function_call(context: IRContext):
            pass

        # 类型注解在 type_analysis 阶段的 pre 时机执行
        @extension_processor(
            "type_alias",
            match_rule=[{"annotation_name": "mcpy_type_alias"}],
            priority=1,
            stage="type_analysis",
            timing="pre"
        )
        def process_type_alias(context: IRContext):
            pass
    """

    def decorator(func: Callable[[IRContext], None]):
        # 创建动态的ExtensionProcessor子类
        class DynamicExtensionProcessor(ExtensionProcessor):
            def __init__(self):
                # 先设置属性，再调用父类初始化（父类会调用_setup_rules）
                self.process_func = func
                self.priority = priority
                self.match_rule = match_rule
                self.stage = stage
                self.timing = timing
                self.original_func: Optional[Callable] = None
                super().__init__(name)

            def _setup_rules(self) -> None:
                """设置匹配规则"""
                # 处理 match_rule 参数
                if self.match_rule is not None:
                    self.add_rule(
                        MatchRule(conditions=self.match_rule, priority=self.priority)
                    )

            def process(self, context: IRContext) -> None:
                """调用用户定义的处理函数"""
                self.process_func(context)

        # 设置类名和模块信息
        class_name = f"{name.title().replace('_', '')}Extension"
        DynamicExtensionProcessor.__name__ = class_name
        DynamicExtensionProcessor.__qualname__ = class_name

        # 创建实例并返回
        extension_instance = DynamicExtensionProcessor()

        # 将原函数作为属性保存，方便调试和测试
        extension_instance.original_func = func

        register_extension(extension_instance)
        return extension_instance

    return decorator
